import os
import cv2
import json
from tqdm import tqdm


def load_spire_dir(spire_dir):
    """
    :param spire_dir: path to spire dataset root dir
    :return: List - img_list, List - ann_list
    """
    img_list = []
    ann_list = []
    for sub_dir in os.listdir(spire_dir):
        sub_fn = os.path.join(spire_dir, sub_dir)
        if os.path.isdir(sub_fn):
            if sub_dir == 'scaled_images':
                img_list.extend([os.path.join(sub_fn, img_nm) for img_nm in os.listdir(sub_fn) if
                                 img_nm.endswith('.jpg') or img_nm.endswith('.png')])
                ann_list.extend([os.path.join(spire_dir, 'annotations', img_nm + '.json') for img_nm in
                                 os.listdir(sub_fn) if
                                 img_nm.endswith('.jpg') or img_nm.endswith('.png')])
            else:
                for sub_sub_dir in os.listdir(sub_fn):
                    sub_sub_fn = os.path.join(sub_fn, sub_sub_dir)
                    if os.path.isdir(sub_sub_fn):
                        if sub_sub_dir == 'scaled_images':
                            img_list.extend([os.path.join(sub_sub_fn, img_nm) for img_nm in os.listdir(sub_sub_fn) if
                                             img_nm.endswith('.jpg') or img_nm.endswith('.png')])
                            ann_list.extend([os.path.join(sub_fn, 'annotations', img_nm + '.json') for img_nm in
                                             os.listdir(sub_sub_fn) if
                                             img_nm.endswith('.jpg') or img_nm.endswith('.png')])
    return img_list, ann_list


def format_check(spire_root_dir):
    img_list, ann_list = load_spire_dir(spire_root_dir)
    all_cates = []

    pre_cates = [['dji-phantom', 'unknown-multi-rotor', 'dji-m210', 'bit-fl5',
                  'unknown-multi-rotor-infrared', 'dji-mavic', 'dji-inspire',
                  'dji-tello', 'dji-m100', 'bit-lf1', 'xinhua-m1'],
                 ['unknown-fixed-wing', 'unknown-fixed-wing-infrared', 'switchblade-300', 'liuhe-lh2300'],
                 ['unknown-bird'], ['unknown-moving-object']]
    cgd_cates = ['multi_rotor', 'fixed_wing', 'bird', 'moving_object']
    for i in tqdm(range(len(img_list))):
        img_fn = img_list[i]
        ann_fn = ann_list[i]
        img = cv2.imread(img_fn)
        ann = json.loads(open(ann_fn, 'r').read())
        for j in range(len(ann['annos'])):
            for k, pre_cate in enumerate(pre_cates):
                if ann['annos'][j]['category_name'] in pre_cate:
                    ann['annos'][j]['category_name'] = cgd_cates[k]
        fp = open(ann_fn, 'w')
        json.dump(ann, fp)
        fp.close()

    print('Check Done! Success!')
    print(all_cates)


if __name__ == '__main__':
    spire_root = r'C:\dataset\BB-Air2AirDet-trainval-v230503'
    format_check(spire_root)

